import sqlite3
from multiprocessing import Pool
import json
import baostock as bs
import pandas as pd
from concurrent.futures import ThreadPoolExecutor
import time

def timeit_test(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        func(*args, **kwargs)
        elapsed = (time.perf_counter() - start)
        print('Time used:{}'.format(elapsed))
    return wrapper

def bs_k_data_stock(code_val='sz.000651', start_val='2009-01-01', end_val='2020-10-28',
                freq_val='d', adjust_val='3'):
    bs.login()
    # 获取历史行情数据
    fields = "date,open,high,low,close,volume"
    df_bs = bs.query_history_k_data(code_val, fields, start_date=start_val,
                end_date=end_val, frequency=freq_val, adjustflag=adjust_val)
    # frequency="d" 取日k线， adjustflag="3"默认不复权 1::后复权 2:前复权

    data_list=[]
    while(df_bs.error_code == '0') & df_bs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(df_bs.get_row_data())

    result = pd.DataFrame(data_list, columns=df_bs.fields)
    result.close = result.close.astype('float64')
    result.open = result.open.astype('float64')
    result.low = result.low.astype('float64')
    result.high = result.high.astype('float64')
    result.volume = result.volume.astype('int')
    result.volume = result.volume/100 # 单位转换: 股-手
    
    result.date = pd.DatetimeIndex(result.date)
    result.set_index("date", drop=True, inplace=True)
    result.index = result.index.set_names('Date')

    recon_data = {'Date':result.index,
            'Open': result.open, 
            'Close':result.close,
            'High':result.high,
            'Low':result.low,
            'Volume':result.volume}
    df_recon = pd.DataFrame(recon_data)
    # 退出系统
    bs.logout()
    return df_recon

def json_to_str():
    # load 将文件中的字符串变换为数据类型
    with open("/home/colby/sd_480/01.liang_hua_tou_zi/01.python_gu_piao_liang_hua/00.huo_qu_gu_piao_shu_ju/stock_pool_baostock.json", 'r', encoding='utf-8') as load_f:
        stock_index = json.load(load_f)

    return stock_index

def map_fun( code, start='2018-01-01', end='2020-11-06'):
    try:
        df_data = bs_k_data_stock(code_val=code, start_val=start, end_val=end)
        df_data.to_sql(""+code, conn, index=False, if_exists='append')
        print("right code is %s" % code)
    except:
        print("error code is %s" % code)
    # df = pd.read_sql_query("select * from  \"" + code + "\"", conn)
    # print(df)

# @timeit_test
def stock_to_sql( start='2009-01-01', end='2020-10-28'):
    stock_index = json_to_str() # 读取股票池 Json 文件
    itr_arg = [code for code in stock_index['股票'].values()]
    pool = Pool(1)
    pool.map(map_fun, itr_arg[:])
    pool.close() # 关闭进程池，不再接受新的进程
    pool.join() # 主进程阻塞等待 子进程的退出

if __name__ == "__main__":

    conn = sqlite3.connect('/home/colby/sd_480/01.liang_hua_tou_zi/01.python_gu_piao_liang_hua/00.huo_qu_gu_piao_shu_ju/stock-data.db')
    c = conn.cursor()
    stock_to_sql()
    conn.commit()# 提交对数据库的改动
    conn.close() # 关闭数据库